#include "vio/front_end/feature_detection.h"

namespace vio {
namespace feature_detection {
AbstractDetector::AbstractDetector(const int img_width,
                                   const int img_height,
                                   const int cell_size,
                                   const int n_pyr_levels)
    : cell_size_(cell_size),
      n_pyr_levels_(n_pyr_levels),
      grid_n_cols_(std::ceil(static_cast<double>(img_width) / cell_size_)),
      grid_n_rows_(std::ceil(static_cast<double>(img_height) / cell_size_)),
      grid_occupancy_(grid_n_cols_ * grid_n_rows_, false) {}

void AbstractDetector::ResetGrid() {
  std::fill(grid_occupancy_.begin(), grid_occupancy_.end(), false);
}

void AbstractDetector::SetExistingFeatures(const Features& fts) {
  std::for_each(fts.begin(), fts.end(), [&](Feature::Ptr fi) {
    grid_occupancy_.at(static_cast<int>(fi->px[1] / cell_size_) * grid_n_cols_ +
                       static_cast<int>(fi->px[0] / cell_size_)) = true;
  });
}

void AbstractDetector::SetGridOccupancy(const Eigen::Vector2d& px) {
  grid_occupancy_.at(static_cast<int>(px[1] / cell_size_) * grid_n_cols_ +
                     static_cast<int>(px[0] / cell_size_)) = true;
}

FastDetector::FastDetector(const int img_width,
                           const int img_height,
                           const int cell_size,
                           const int n_pyr_levels)
    : AbstractDetector(img_width, img_height, cell_size, n_pyr_levels) {}

void FastDetector::Detect(Frame::Ptr frame,
                          const ImgPyr& img_pyr,
                          const double detection_threshold,
                          Features& fts) {
  Corners corners(grid_n_cols_ * grid_n_rows_,
                  Corner(0, 0, detection_threshold, 0, 0.0f));
  for (int L = 0; L < n_pyr_levels_; L++) {
    const int scale = (1 << L);
    std::vector<fast::fast_xy> fast_corners;
#if __SSE2__
    fast::fast_corner_detect_10_sse2((fast::fast_byte*)img_pyr[L].data,
                                     img_pyr[L].cols,
                                     img_pyr[L].rows,
                                     img_pyr[L].cols,
                                     20,
                                     fast_corners);
#else
    fast::fast_corner_detect_10((fast::fast_byte*)img_pyr[L].data,
                                img_pyr[L].cols,
                                img_pyr[L].rows,
                                img_pyr[L].cols,
                                20,
                                fast_corners);
#endif
    std::vector<int> scores, num_corners;
    fast::fast_corner_score_10((fast::fast_byte*)img_pyr[L].data,
                               img_pyr[L].cols,
                               fast_corners,
                               20,
                               scores);
    fast::fast_nonmax_3x3(fast_corners, scores, num_corners);

    for (auto it = num_corners.begin(), ite = num_corners.end(); it != ite;
         it++) {
      fast::fast_xy& xy = fast_corners.at(*it);
      const int k =
          static_cast<int>((xy.y * scale) / cell_size_) * grid_n_cols_ +
          static_cast<int>((xy.x * scale) / cell_size_);
      if (grid_occupancy_[k]) {
        continue;
      }
      const float score = vision_utils::ShiTomasiScore(img_pyr[L], xy.x, xy.y);
      if (score > corners.at(k).score) {
        corners.at(k) = Corner(xy.x * scale, xy.y * scale, score, L, 0.0f);
      }
    }
  }
  // create feature for every corner that has high enough corner score
  std::for_each(corners.begin(), corners.end(), [&](Corner& c) {
    if (c.score > detection_threshold) {
      Feature::Ptr f_ptr =
          std::make_shared<Feature>(frame, Eigen::Vector2d(c.x, c.y), c.level);
      fts.push_back(f_ptr);
    }
  });

  ResetGrid();
}

}  // namespace feature_detection
}  // namespace vio